{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Video Understanding\n",
    "\n",
    "In this example, we will generate a description of a video using `Qwen2-VL`, `Qwen2-5-VL`, `LLava`, and `Idefics3`, with more models coming soon.\n",
    "\n",
    "This feature is currently in beta, may not work as expected.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U mlx-vlm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from mlx_vlm import load, generate\n",
    "from mlx_vlm.video_generate import process_vision_info\n",
    "\n",
    "import mlx.core as mx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model and processor\n",
    "model, processor = load(\"mlx-community/Qwen2.5-VL-7B-Instruct-4bit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Messages containing a video and a text query\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"video\",\n",
    "                \"video\": \"videos/fastmlx_local_ai_hub.mp4\",\n",
    "                \"max_pixels\": 360 * 360,\n",
    "                \"fps\": 1.0,\n",
    "            },\n",
    "            {\"type\": \"text\", \"text\": \"Describe this video.\"},\n",
    "        ],\n",
    "    }\n",
    "]\n",
    "\n",
    "# Preparation for inference\n",
    "text = processor.apply_chat_template(\n",
    "    messages, tokenize=False, add_generation_prompt=True\n",
    ")\n",
    "image_inputs, video_inputs = process_vision_info(messages)\n",
    "inputs = processor(\n",
    "    text=[text],\n",
    "    images=image_inputs,\n",
    "    videos=video_inputs,\n",
    "    padding=True,\n",
    "    return_tensors=\"pt\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert inputs to mlx arrays\n",
    "input_ids = mx.array(inputs['input_ids'])\n",
    "pixel_values = mx.array(inputs['pixel_values_videos'])\n",
    "mask = mx.array(inputs['attention_mask'])\n",
    "video_grid_thw = mx.array(inputs['video_grid_thw'])\n",
    "\n",
    "kwargs = {\n",
    "    \"video_grid_thw\": video_grid_thw,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kwargs[\"video\"] = \"videos/fastmlx_local_ai_hub.mp4\"\n",
    "kwargs[\"input_ids\"] = input_ids\n",
    "kwargs[\"pixel_values\"] = pixel_values\n",
    "kwargs[\"mask\"] = mask\n",
    "response = generate(model, processor, prompt=text, temperature=0.7, max_tokens=100, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open video and play it\n",
    "from ipywidgets import Video\n",
    "Video.from_file(\"videos/fastmlx_local_ai_hub.mp4\", width=320, height=240)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
